# -- coding:UTF-8 --
import pstats
import numpy as np
import cv2
import random
import time
import matplotlib.pyplot as plt
import pyrealsense2 as rs
import cv2
import numpy as np

"""
Created on  Oct 11 2020
静态版vibe
@author: Kkoutianwu
"""


class ViBe:
    '''0
    ViBe运动检测，分割背景和前景运动图像
    '''

    def __init__(self, num_sam=4, radius=15, num_frames=0, count_background=3):
        self.num_sam = num_sam  # 每个像素的样本集数量，默认5个
        self.radius = radius  # 匹配半径，即在该半径内则认为是匹配像素
        self.num_frames = num_frames  # 用于初始化动态阈值的帧数
        self.count_background = count_background  # 每个像素点被认为是背景的次数大于该值就更新其样本集和动态阈值
        self.background = 0
        self.foreground = 255

    def ProcessFirstFrame(self, img):
        '''
        处理视频的第一帧
        1、初始化每个像素的样本集矩阵
        2、初始化前景矩阵的mask
        3、初始化前景像素的检测次数矩阵
        参数：
        img: 传入的numpy图像素组，要求RGB图像
        返回值：
        每个像素的样本集numpy数组
        '''

        self.samples = np.zeros((self.num_sam, img.shape[0], img.shape[1]), dtype=np.uint8)
        self.thresholds = np.ones((img.shape[0], img.shape[1]), dtype=np.uint8)*50

        self.fgCount = np.zeros((img.shape[0], img.shape[1]))  # 每个像素被检测为前景的次数
        self.fgMask = np.zeros((img.shape[0], img.shape[1]),dtype=np.uint8)  # 保存前景像素
        self.foreground_area = np.zeros((img.shape[0], img.shape[1]), dtype=np.uint8)
        self.count = 0  # 当前帧数
        self.__buildNeighborArray(img)

    def __buildNeighborArray(self, img):
        '''
        构建一副图像中每个像素的邻域数组
        参数：输入灰度图像
        返回值：每个像素9邻域数组，保存到self.samples中
        '''

        height, width, _ = img.shape
        img_b, img_g, img_r = cv2.split(img)
        contrast = cv2.absdiff(img_b, img_g) + cv2.absdiff(img_g, img_r) + cv2.absdiff(img_b, img_r)  # 计算对比度
        contrast[contrast > 255] = 255

        img = cv2.merge([img_b, img_g, img_r, contrast])

        # 生成随机偏移数组，用于计算随机选择的邻域坐标
        ramoff_xy = np.random.randint(-3, 4, size=(2, self.num_sam, height, width))
        xr_ = np.tile(np.arange(width), (height, 1))
        print ("xr_.shape", xr_.shape)
        yr_ = np.tile(np.arange(height), (width, 1)).T
        print ("yr_.shape", yr_.shape)

        xyr_ = np.zeros((2, self.num_sam, height, width))
        for i in range(self.num_sam):
            xyr_[1, i] = xr_
            xyr_[0, i] = yr_

        xyr_ = xyr_ + ramoff_xy

        xyr_[xyr_ < 0] = 0
        tpr_ = xyr_[1, :, :, :]
        tpr_[tpr_ >= width] = width - 1
        tpb_ = xyr_[0, :, :, :]
        tpb_[tpb_ >= height] = height - 1
        xyr_[0, :, :, :] = tpb_
        xyr_[1, :, :, :] = tpr_
        xyr = xyr_.astype(int)

        self.samples = img[xyr[0, :, :, :], xyr[1, :, :, :]]



    def Diff(self, img):
        '''
        处理每帧视频，更新运动前景，并更新样本集。该函数是本类的主函数
        输入：灰度图像
        '''

        height, width, _ = img.shape
        img_b, img_g, img_r = cv2.split(img)

        contrast = cv2.absdiff(img_b, img_g) + cv2.absdiff(img_g, img_r) + cv2.absdiff(img_b, img_r)
        contrast[contrast > 255] = 255
        self.cur_img = cv2.merge([img_b, img_g, img_r, contrast])

        # # 使用欧式距离
        # dist = np.abs((self.samples.astype(float) - self.cur_img.astype(float)).astype(np.uint16))
        # dist=pow(dist, 2)
        # dist=np.sum(dist,3)
        # dist=np.sqrt(dist)
        # dist[dist > 255] = 255
        # dist=dist.astype(np.uint8)

        # 使用最大差异值
        dist = np.abs((self.samples.astype(float) - self.cur_img.astype(float)).astype(int))
        dist = np.max(dist, axis=3)

        dist = np.min(dist, axis=0)
        self.diff = cv2.convertScaleAbs(dist)
        # 前几帧用于初始化动态阈值
        # if self.count < self.num_frames:
        #     if self.count == 0:
        #         self.thresholds = self.diff
        #     # self.thresholds=(max+self.thresholds)/2
        #     self.thresholds = np.maximum(self.diff, self.thresholds)
        #     if self.count == self.num_frames - 1:
        #         self.thresholds = self.thresholds + self.radius
        #     self.count = self.count + 1
        # print self.count

        # 二值化
        matches = self.diff > self.thresholds
        self.fgMask[matches] = 1
        self.fgMask[~matches] = self.background
        kernel = np.ones((3, 3))
        temp=cv2.filter2D(self.fgMask,-1,kernel)
        self.fgMask[temp>4]=self.foreground
        self.fgMask[temp <= 4] = self.background



        # 背景像素计数+1,前景像素的计数设置为0
        self.fgCount[~matches] = self.fgCount[~matches] + 1
        self.fgCount[matches] = 0
        self.fgCount[self.fgCount > self.count_background] = self.count_background
        return self.fgMask

    def Update(self):
        if self.count >= self.num_frames:
            deny_update_area = (self.foreground_area > 0) | (self.fgCount < self.count_background) # 拒绝更新的区域
            upSelfSamplesInd = np.where(deny_update_area == False)  # 满足可以更新样本集像素的索引

            if self.count == self.num_frames:
                self.pre_pre_diff = self.diff
            elif self.count == self.num_frames+1:
                self.pre_diff=self.diff
            else:
                if self.count % 2 == 0:
                    # 更新动态阈值
                    temp=np.maximum(self.pre_diff,self.pre_pre_diff)
                    temp = np.maximum(temp, self.diff)+self.radius
                    self.pre_pre_diff=self.pre_diff
                    self.pre_diff=self.diff
                    # self.temp = self.thresholds.copy()
                    # # self.temp[self.diff > self.temp] = self.diff[self.diff > self.temp] + self.radius
                    # self.temp=(self.diff+self.temp+10)/2
                    self.thresholds[upSelfSamplesInd] = temp[upSelfSamplesInd]
                    # cv2.imshow("self.thresholds", self.thresholds)

                else:
                    pass
                    # # 更新背景样本集
                    # upSelfSamplesPosition = np.random.randint(self.num_sam, size=upSelfSamplesInd[0].shape)  # 生成随机更新自己样本集的的索引
                    # samInd = (upSelfSamplesPosition, upSelfSamplesInd[0], upSelfSamplesInd[1])
                    # self.samples[samInd] = self.cur_img[upSelfSamplesInd]  # 更新自己样本集中的一个样本为本次图像中对应像素值
            self.count+=1

def main():
    vibe = ViBe()
    pipeline = rs.pipeline()
    config = rs.config()
    config.enable_stream(rs.stream.color, 1280, 720, rs.format.bgr8, 30)
    pipeline.start(config)
    while True:
        key = cv2.waitKey(30)
        frames = pipeline.wait_for_frames()
        color_frame = frames.get_color_frame()
        img1 = np.asanyarray(color_frame.get_data())
        img1 = cv2.GaussianBlur(img1, (5, 5), 0, 0)
        cv2.imshow("i", img1)
        if key == 27:
            vibe.ProcessFirstFrame(img1)
            break
    while True:
        frames = pipeline.wait_for_frames()
        color_frame = frames.get_color_frame()
        img1 = np.asanyarray(color_frame.get_data())
        # img1 = cv2.GaussianBlur(img1, (5, 5), 0, 0)
        cv2.imshow("i", img1)
        # 输出二值图
        # (segMat, samples) = update(gray, samples)
        segMat=vibe.Diff(img1)
        # 形态学处理模板初始化
        kernel_offset = np.array((
            [1, 2, 2, 2, 1],
            [2, -1, -1, -1, 2],
            [2, -1, -10, -1, 2],
            [2, -1, -1, -1, 2],
            [1, 2, 2, 2, 1]), dtype="float32") / 5
        weights_offset = cv2.filter2D(segMat.astype(np.double), -1, kernel_offset)
        weights_offset[weights_offset < 0] = 0
        # cv2.imshow("weights_offset", cv2.convertScaleAbs(weights_offset))
        # weights_offset = cv2.threshold(weights_offset, threshold_value, 255, cv2.THRESH_BINARY)[1]
        weights_offset = (weights_offset.astype(np.double) / 50)

        result = cv2.convertScaleAbs(segMat.astype(np.double) * weights_offset)
        # result[result > threshold_value] = 255
        result = cv2.threshold(result, 250, 255, cv2.THRESH_BINARY)[1]
        cv2.imshow("result ", result)

        # 寻找轮廓
        contours, hierarchy = cv2.findContours(result, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        vibe.foreground_area[:, :] = 0
        for i in range(0, len(contours)):
            x, y, w, h = cv2.boundingRect(contours[i])
            if w * h > 1000:
                # print(x,y,w,h)
                cv2.rectangle(img1, (x, y), (x + w, y + h), (0, 255, 0), 2)
                vibe.foreground_area[y:y + h, x: x + w] = 255
        cv2.imshow("frame", img1)
        cv2.imshow("SegMat", segMat)
        vibe.Update()
        # cv2.imwrite("./result/" + str(c) + ".jpg", frame,[int(cv2.IMWRITE_PNG_STRATEGY)])
        k = cv2.waitKey(1)
        if k == 27:
            cv2.destroyAllWindows()
            break


if __name__ == '__main__':
    main()
